import os
import traceback

from openai import OpenAI
from tenacity import before_sleep_log, retry, retry_if_exception_type, stop_after_attempt, wait_exponential

from src import config
from src.utils import get_docker_safe_url, logger


class OpenAIBase:
    def __init__(self, api_key, base_url, model_name, **kwargs):
        self.api_key = api_key
        self.base_url = base_url
        self.client = OpenAI(api_key=api_key, base_url=base_url)
        self.model_name = model_name
        self.info = kwargs

    @retry(
        stop=stop_after_attempt(3),
        wait=wait_exponential(multiplier=1, min=1, max=10),
        retry=retry_if_exception_type((Exception,)),
        before_sleep=before_sleep_log(logger, log_level="WARNING"),
        reraise=True,
    )
    def call(self, message, stream=False):
        if isinstance(message, str):
            messages = [{"role": "user", "content": message}]
        else:
            messages = message

        try:
            if stream:
                response = self._stream_response(messages)
            else:
                response = self._get_response(messages)

        except Exception as e:
            err = (
                f"Error streaming response: {e}, URL: {self.base_url}, "
                f"API Key: {self.api_key[:5]}***, Model: {self.model_name}"
            )
            logger.error(err)
            raise Exception(err)

        return response

    def _stream_response(self, messages):
        response = self.client.chat.completions.create(
            model=self.model_name,
            messages=messages,
            stream=True,
        )
        for chunk in response:
            if len(chunk.choices) > 0:
                yield chunk.choices[0].delta

    def _get_response(self, messages):
        response = self.client.chat.completions.create(
            model=self.model_name,
            messages=messages,
            stream=False,
        )
        return response.choices[0].message

    def get_models(self):
        try:
            return self.client.models.list(extra_query={"type": "text"})
        except Exception as e:
            logger.error(f"Error getting models: {e}")
            return []


class OpenModel(OpenAIBase):
    def __init__(self, model_name=None):
        model_name = model_name or "gpt-4o-mini"
        api_key = os.getenv("OPENAI_API_KEY")
        base_url = os.getenv("OPENAI_API_BASE")
        super().__init__(api_key=api_key, base_url=base_url, model_name=model_name)


class VLLMModel(OpenAIBase):
    """
    通用 vLLM 服务封装（OpenAI-compatible）
    自动适配任意模型，只要 vLLM 服务在 base_url 提供 /v1/chat/completions 接口即可。
    """
    def __init__(self, model_name=None):
        model_name = model_name or os.getenv("VLLM_DEFAULT_MODEL", "default")
        api_key = os.getenv("VLLM_API_KEY", "none")
        base_url = os.getenv("VLLM_API_BASE", "http://localhost:8020/v1")

        logger.info(f"> Using vLLM model `{model_name}` from {base_url}")
        super().__init__(api_key=api_key, base_url=base_url, model_name=model_name)


class CustomModel(OpenAIBase):
    def __init__(self, model_info):
        model_name = model_info["name"]
        api_key = model_info.get("api_key") or "custom_model"
        base_url = get_docker_safe_url(model_info["api_base"])
        logger.info(f"> Custom model: {model_name}, base_url: {base_url}")

        super().__init__(api_key=api_key, base_url=base_url, model_name=model_name)


class GeneralResponse:
    def __init__(self, content):
        self.content = content
        self.is_full = False


def select_model(model_provider, model_name=None):
    """根据模型提供者选择模型"""
    assert model_provider is not None, "Model provider not specified"
    model_info = config.model_names.get(model_provider, {})
    model_name = model_name or model_info.get("default", "")

    logger.info(f"Selecting model from `{model_provider}` with `{model_name}`")

    # === OpenAI ===
    if model_provider == "openai":
        return OpenModel(model_name)

    # === vLLM ===
    if model_provider == "vllm":
        return VLLMModel(model_name)

    # === Custom ===
    if model_provider == "custom":
        model_info = get_custom_model(model_name)
        return CustomModel(model_info)

    # === Other ===
    try:
        model = OpenAIBase(
            api_key=os.getenv(model_info["env"][0]),
            base_url=model_info["base_url"],
            model_name=model_name,
        )
        return model
    except Exception as e:
        raise ValueError(f"Model provider {model_provider} load failed, {e} \n {traceback.format_exc()}")


def get_custom_model(model_id):
    """return model_info"""
    assert config.custom_models is not None, "custom_models is not set"
    modle_info = next((x for x in config.custom_models if x["custom_id"] == model_id), None)
    if modle_info is None:
        raise ValueError(f"Model {model_id} not found in custom models")
    return modle_info


async def test_chat_model_status(provider: str, model_name: str) -> dict:
    """
    测试指定聊天模型的状态

    Args:
        provider: 模型提供商
        model_name: 模型名称

    Returns:
        dict: 包含状态信息的字典
    """
    try:
        # 加载模型
        logger.debug(f"Selecting chat model {provider}/{model_name}")
        model = select_model(provider, model_name)

        # 使用简单的测试消息
        test_messages = [{"role": "user", "content": "Say 1"}]

        # 发送测试请求
        response = model.call(test_messages, stream=False)
        logger.debug(f"Test chat model status response: {response}")

        # 检查响应是否有效
        if response and response.content:
            return {"provider": provider, "model_name": model_name, "status": "available", "message": "连接正常"}
        else:
            return {"provider": provider, "model_name": model_name, "status": "unavailable", "message": "响应无效"}

    except Exception as e:
        error_msg = str(e)
        # 尝试提取更友好的错误信息
        if "503" in error_msg and "too busy" in error_msg.lower():
            error_msg = "API 服务繁忙，请稍后重试"
        elif "Error code:" in error_msg:
            # 提取关键错误信息
            try:
                if "'message':" in error_msg:
                    import re
                    match = re.search(r"'message':\s*'([^']+)'", error_msg)
                    if match:
                        error_msg = match.group(1)
            except Exception:
                pass
        
        logger.error(f"测试聊天模型状态失败 {provider}/{model_name}: {error_msg}")
        return {"provider": provider, "model_name": model_name, "status": "error", "message": error_msg}


async def test_all_chat_models_status() -> dict:
    """
    测试所有支持的聊天模型状态

    Returns:
        dict: 包含所有模型状态的字典
    """
    from src import config

    results = {}

    # 获取所有可用的模型
    for provider, provider_info in config.model_names.items():
        if provider == "custom":
            # 处理自定义模型
            for custom_model in config.custom_models:
                model_id = f"custom/{custom_model['custom_id']}"
                status = await test_chat_model_status("custom", custom_model["custom_id"])
                results[model_id] = status
        else:
            # 处理普通模型
            for model_name in provider_info.models:
                model_id = f"{provider}/{model_name}"
                status = await test_chat_model_status(provider, model_name)
                results[model_id] = status

    available_count = len([m for m in results.values() if m["status"] == "available"])

    return {"models": results, "total": len(results), "available": available_count}


if __name__ == "__main__":
    pass
